#include <math.h>
#include <stdio.h>
#include <string.h>
#include <algorithm>
#include <cstdint>
#include <string>
#include <vector>
#include <algorithm>
#include "typeConv.hpp"
#include "common.h"
#include "yolox_post_process.h"
#ifdef LYNXI_PLUGIN
#include "lyn_plugin_dev.h"
#endif

using namespace COMMON;

struct Yolov3Config {
  int num_anchor;
  std::vector<int> strides;
  float anchorsall[18];  // anchors 比例
  std::vector<std::string> class_names;
};

static const Yolov3Config default_yolov3_config = {
    3,
    {32, 16, 8},
    {10.f, 13.f, 16.f, 30.f, 33.f, 23.f, 30.f, 61.f, 62.f, 45.f, 59.f, 119.f,
     116.f, 90.f, 156.f, 198.f, 373.f, 326.f},
    {"person",        "bicycle",      "car",
     "motorcycle",    "airplane",     "bus",
     "train",         "truck",        "boat",
     "traffic light", "fire hydrant", "stop sign",
     "parking meter", "bench",        "bird",
     "cat",           "dog",          "horse",
     "sheep",         "cow",          "elephant",
     "bear",          "zebra",        "giraffe",
     "backpack",      "umbrella",     "handbag",
     "tie",           "suitcase",     "frisbee",
     "skis",          "snowboard",    "sports ball",
     "kite",          "baseball bat", "baseball glove",
     "skateboard",    "surfboard",    "tennis racket",
     "bottle",        "wine glass",   "cup",
     "fork",          "knife",        "spoon",
     "bowl",          "banana",       "apple",
     "sandwich",      "orange",       "broccoli",
     "carrot",        "hot dog",      "pizza",
     "donut",         "cake",         "chair",
     "couch",         "potted plant", "bed",
     "dining table",  "toilet",       "tv",
     "laptop",        "mouse",        "remote",
     "keyboard",      "cell phone",   "microwave",
     "oven",          "toaster",      "sink",
     "refrigerator",  "book",         "clock",
     "vase",          "scissors",     "teddy bear",
     "hair drier",    "toothbrush"}};

static inline float sigmoid(float x) {
  return static_cast<float>(1.f / (1.f + exp(-x)));
}

static inline float arcsigmoid(float x) {
  return static_cast<float>(log(x / (1.f - x)));
}

static inline float intersection_area(const BboxResult &a,
                                      const BboxResult &b) {
  if (a.xmin > b.xmax || a.xmax < b.xmin || a.ymin > b.ymax ||
      a.ymax < b.ymin) {
    // no intersection
    return 0.f;
  }

  float inter_width = std::min(a.xmax, b.xmax) - std::max(a.xmin, b.xmin);
  float inter_height = std::min(a.ymax, b.ymax) - std::max(a.ymin, b.ymin);

  return inter_width * inter_height;
}

static inline void qsort_descent_inplace(std::vector<BboxResult> &datas,
                                         int left, int right) {
  int i = left;
  int j = right;
  float p = datas[(left + right) / 2].score;

  while (i <= j) {
    while (datas[i].score > p) i++;

    while (datas[j].score < p) j--;

    if (i <= j) {
      // swap
      std::swap(datas[i], datas[j]);

      i++;
      j--;
    }
  }

  if (left < j) qsort_descent_inplace(datas, left, j);

  if (i < right) qsort_descent_inplace(datas, i, right);
}

static void qsort_descent_inplace(std::vector<BboxResult> &datas) {
  if (datas.empty()) return;

  qsort_descent_inplace(datas, 0, static_cast<int>(datas.size() - 1));
}

static void nms_sorted_bboxes(const std::vector<BboxResult> &bboxes,
                              std::vector<size_t> &picked,
                              float nms_threshold) {
  picked.clear();

  const size_t n = bboxes.size();

  for (size_t i = 0; i < n; i++) {
    const BboxResult &a = bboxes[i];

    int keep = 1;
    for (int j = 0; j < (int)picked.size(); j++) {
      const BboxResult &b = bboxes[picked[j]];

      // intersection over union
      float inter_area = intersection_area(a, b);
      float union_area = a.area + b.area - inter_area;
      //             float IoU = inter_area / union_area
      if (inter_area / union_area > nms_threshold) keep = 0;
    }

    if (keep) picked.push_back(i);
  }
}

static int count_n = 0;

/**
 *  yolov3 输出feature map
 * [1, 13, 13, 255]
 *  x, y , w, h, bbox_conf, cls_conf...
 * 当输出shape为[1, 255, 13, 13]时，
 * 单个box的数据都在channel上，首先找到该像素点对应anchor在channel上的首地址
 * 比如对于anchor n,其对应feature map i, j上像素点对应的channel首地址为
 * n*per_anchor_channels*feature_h*feature_w + i*feature_w + j
 */
//// tensor 解码，边框回归，输入是 tensor
/// data的指针及相应配参，输出是检测结果结构体vector
// [1, 13, 13, 255]
void decode_tensor_transpose(const Yolov3Config &yolov3_config, std::vector<BboxResult> &dets,
                             const int16_t *out_vector, const float *anchors,
                             float obj_thresh, int16_t box_score_threshold,
                             const YoloxPostProcessInfo_t *post_info, int grid_h,
                             int grid_w) {
  int class_num = post_info->class_num;
  int b_chan = 5 + class_num;
  int nb_chan = yolov3_config.num_anchor * b_chan;

  int stride = post_info->height / grid_h;

  for (int i = 0; i < grid_h; ++i) {
    for (int j = 0; j < grid_w; ++j) {
      const int16_t *bbox_pred_ptr =
          out_vector + i * grid_w * nb_chan + j * nb_chan;

      for (int na = 0; na < yolov3_config.num_anchor; ++na) {
        const float anchor_w = anchors[2 * na];
        const float anchor_h = anchors[2 * na + 1];

        const int16_t *pred_ptr = bbox_pred_ptr + na * b_chan;

        int16_t objness = pred_ptr[4];

        // if (count_n == 2 && i >= 25 && j >= 25) {
        //   float mmxxx = half2float(objness);
        //   if (mmxxx > -2.f)
        //     LOG_PLUGIN_E("%f\n", mmxxx);
        // }

        // find class index with max class score
        int class_index = 0;
        int16_t class_score = 0;
        if (class_num == 80) {
          for (int q = 0; q < class_num; q++) {
            int16_t score = pred_ptr[q + 5];
            if (score > class_score) {
              class_index = q;
              class_score = score;
            }
          }
        } else {
          class_index = pred_ptr[5];
          class_score = pred_ptr[6];
        }

        float bboxScore =
            sigmoid(half2float(objness)) * sigmoid(half2float(class_score));
        if (bboxScore < obj_thresh) {
          continue;
        }

        BboxResult box;

        float x = (j + sigmoid(half2float(pred_ptr[0]))) * stride;
        float y = (i + sigmoid(half2float(pred_ptr[1]))) * stride;
        float w = anchor_w * static_cast<float>(exp(half2float(pred_ptr[2])));
        float h = anchor_h * static_cast<float>(exp(half2float(pred_ptr[3])));

        box.area = h * w;
        box.xmin = x - w * 0.5f;
        box.ymin = y - h * 0.5f;
        box.xmax = x + w * 0.5f;
        box.ymax = y + h * 0.5f;

        box.id = class_index;
        box.score = bboxScore;
        strncpy(box.class_name, yolov3_config.class_names[class_index].c_str(), sizeof(box.class_name) - 1);
        dets.emplace_back(box);
      }
    }
  }

  count_n++;
}

// [1, 3, 13*13, 85]
// [:, :, :, 0:2].sigmoid
// [:, :, :, 2:4].exp
// [:, :, :, 4:].sigmoid
void decode_tensor_fast(const Yolov3Config &yolov3_config, std::vector<BboxResult> &dets,
                        const int16_t *out_vector, const float *anchors,
                        float obj_thresh, int16_t box_score_threshold,
                        const YoloxPostProcessInfo_t *post_info, int grid_h,
                        int grid_w) {
  int class_num = post_info->class_num;
  int b_chan = 5 + class_num;
  // int nb_chan = yolov3_config.num_anchor * b_chan;
  int gird_stride = grid_h * grid_w;

  int stride = post_info->height / grid_h;

  for (int na = 0; na < yolov3_config.num_anchor; ++na) {
    const float anchor_w = anchors[2 * na];
    const float anchor_h = anchors[2 * na + 1];

    for (int i = 0; i < grid_h; ++i) {
      for (int j = 0; j < grid_w; ++j) {
        const int16_t *pred_ptr = out_vector + na * gird_stride * b_chan +
                                  i * grid_w * b_chan + j * b_chan;

        int16_t objness = pred_ptr[4];
        if (objness < box_score_threshold) {
          continue;
        }

        // find class index with max class score
        int class_index = 0;
        int16_t class_score = 0;
        if (class_num == 80) {
          for (int q = 0; q < class_num; q++) {
            int16_t score = pred_ptr[q + 5];
            if (score > class_score) {
              class_index = q;
              class_score = score;
            }
          }
        } else {
          class_index = pred_ptr[5];
          class_score = pred_ptr[6];
        }

        float bboxScore = half2float(objness) * half2float(class_score);
        if (bboxScore < obj_thresh) {
          continue;
        }

        BboxResult box;

        float x = (j + half2float(pred_ptr[0])) * stride;
        float y = (i + half2float(pred_ptr[1])) * stride;
        float w = anchor_w * half2float(pred_ptr[2]);
        float h = anchor_h * half2float(pred_ptr[3]);

        box.area = h * w;
        box.xmin = x - w * 0.5f;
        box.ymin = y - h * 0.5f;
        box.xmax = x + w * 0.5f;
        box.ymax = y + h * 0.5f;

        box.id = class_index;
        box.score = bboxScore;
        strncpy(box.class_name, yolov3_config.class_names[class_index].c_str(), sizeof(box.class_name) - 1);
        dets.emplace_back(box);
      }
    }
  }

  count_n++;
}

int lynYoloxPostProcess(YoloxPostProcessInfo_t *post_info) {
  int16_t *output_tensor = nullptr;
  lynBoxesInfo *boxesInfo = nullptr;
#ifdef LYNXI_PLUGIN
  output_tensor = (int16_t *)lynPluginGetVirtAddr(post_info->output_tensor);
  if (output_tensor == nullptr) {
    LOG_PLUGIN_E("get output tensor addr error\n");
    return -1;
  }

  boxesInfo = (lynBoxesInfo *)lynPluginGetVirtAddr(post_info->boxesInfo);
  if (boxesInfo == nullptr) {
    LOG_PLUGIN_E("get boxesInfo addr error\n");
    return -1;
  }
#else
  output_tensor = (int16_t *)post_info->output_tensor;
  boxesInfo = post_info->boxesInfo;
#endif
  Yolov3Config yolov3_config = default_yolov3_config;
  generateClassNameAndNum(post_info->labelList,
                          yolov3_config.class_names,
                          post_info->class_num);

  int predict_per_grid =
      (post_info->class_num + 5) * yolov3_config.num_anchor;
  int16_t *tensor0 = (int16_t *)output_tensor;
  int layer0_grid_h = post_info->height / yolov3_config.strides[0];
  int layer0_grid_w = post_info->width / yolov3_config.strides[0];
  const float *layer0_anchor = yolov3_config.anchorsall + 12;

  int16_t *tensor1 = tensor0 + layer0_grid_h * layer0_grid_w * predict_per_grid;
  int layer1_grid_h = post_info->height / yolov3_config.strides[1];
  int layer1_grid_w = post_info->width / yolov3_config.strides[1];
  const float *layer1_anchor = yolov3_config.anchorsall + 6;

  int16_t *tensor2 = tensor1 + layer1_grid_h * layer1_grid_w * predict_per_grid;
  int layer2_grid_h = post_info->height / yolov3_config.strides[2];
  int layer2_grid_w = post_info->width / yolov3_config.strides[2];
  const float *layer2_anchor = yolov3_config.anchorsall;

  int w_pad;
  int h_pad;
  float resize_scale;
  if (post_info->ori_height > post_info->ori_width) {
    resize_scale = (float)post_info->height / post_info->ori_height;
    int new_w = int(post_info->ori_width * resize_scale) & -1;
    w_pad = (post_info->width - new_w) / 2;
    h_pad = 0;
  } else {
    resize_scale = (float)post_info->width / post_info->ori_width;
    int new_h = int(post_info->ori_height * resize_scale) & -1;
    h_pad = (post_info->height - new_h) / 2;
    w_pad = 0;
  }

  std::vector<BboxResult> dets;

  int16_t box_score_threshold = float2half(post_info->score_threshold);

  // 对应输出维度为[1, 3, 13*13, 85]
  // stride 32
  decode_tensor_fast(yolov3_config, dets, tensor0, layer0_anchor, post_info->score_threshold,
                     box_score_threshold, post_info, layer0_grid_h,
                     layer0_grid_w);

  // stride 16
  decode_tensor_fast(yolov3_config, dets, tensor1, layer1_anchor, post_info->score_threshold,
                     box_score_threshold, post_info, layer1_grid_h,
                     layer1_grid_w);

  // stride 8
  decode_tensor_fast(yolov3_config, dets, tensor2, layer2_anchor, post_info->score_threshold,
                     box_score_threshold, post_info, layer2_grid_h,
                     layer2_grid_w);

  // global sort inplace
  qsort_descent_inplace(dets);

  // apply nms
  std::vector<size_t> picked;
  nms_sorted_bboxes(dets, picked, post_info->nms_threshold);

  boxesInfo->boxesNum = std::min((int)picked.size(), BOX_MAX_NUM);

  for (int i = 0; i < boxesInfo->boxesNum; i++) {
    int index = picked[i];

    boxesInfo->boxes[i].xmax = (dets[index].xmax - (float)w_pad) / resize_scale;
    boxesInfo->boxes[i].xmin = (dets[index].xmin - (float)w_pad) / resize_scale;
    boxesInfo->boxes[i].ymax = (dets[index].ymax - (float)h_pad) / resize_scale;
    boxesInfo->boxes[i].ymin = (dets[index].ymin - (float)h_pad) / resize_scale;
    if (static_cast<int>(boxesInfo->boxes[i].xmax) > post_info->ori_width - 1)
      boxesInfo->boxes[i].xmax = post_info->ori_width - 1;
    if (static_cast<int>(boxesInfo->boxes[i].xmin) < 0) boxesInfo->boxes[i].xmin = 0;
    if (static_cast<int>(boxesInfo->boxes[i].ymax) > post_info->ori_height - 1)
      boxesInfo->boxes[i].ymax = post_info->ori_height - 1;
    if (static_cast<int>(boxesInfo->boxes[i].ymin) < 0) boxesInfo->boxes[i].ymin = 0;

    boxesInfo->boxes[i].score = dets[index].score;
    strncpy(boxesInfo->boxes[i].label, dets[index].class_name, sizeof(boxesInfo->boxes[i].label) - 1);
    boxesInfo->boxes[i].id = dets[index].id;
  }

  return 0;
}
